package com.lm.flight.example;

import org.apache.arrow.flight.*;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class FlightClientExample {

    private static final Logger log = LoggerFactory.getLogger(FlightClientExample.class);

    public static void main(String[] args) throws Exception {

        // 创建内存分配器
        try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {

            // 与 Flight Server 建立连接
            Location location = new Location("grpc://localhost:8815");// Location.forGrpcInsecure("localhost", 8815);
            try (FlightClient client = FlightClient.builder(allocator, location).build()) {

                log.info("✅ Connected to Flight server: {}", location);

                // 构造 ticket（与服务端 getStream 的参数匹配）
                Ticket ticket = new Ticket("demo".getBytes());

                // 请求服务端流数据
                try (FlightStream stream = client.getStream(ticket)) {
                    log.info("🚀 Fetching data from Arrow Flight Server...");

                    while (stream.next()) {
                        VectorSchemaRoot root = stream.getRoot();
                        IntVector idVector = (IntVector) root.getVector("id");
                        VarCharVector nameVector = (VarCharVector) root.getVector("name");

                        for (int i = 0; i < root.getRowCount(); i++) {
                            int id = idVector.get(i);
                            String name = new String(nameVector.get(i));
                            log.info("Row {} → id={}, name={}", i, id, name);
                        }
                    }

                    log.info("✅ Data fetch completed.");
                }
            }
        }
    }
}
